#ifndef _SOCKET_HPP
#define _SOCKET_HPP
#include <atomic>
#include <boost/asio/buffer.hpp>
#include <cstddef>
#include <cstdint>
#include <memory>
#include <queue>
#include <iostream>
#include <boost/asio/ip/tcp.hpp>
#include <boost/system/error_code.hpp>
#include "MessageBuffer.hpp"
using boost::asio::ip::tcp;

template <class SocketType>
class Socket : public std::enable_shared_from_this<SocketType> // 因为异步读需要保证回调函数被调用之前，对象是有效的，需要使用shared_from_this
{
public:
    Socket(tcp::socket&& sc) : _socket(std::move(sc)), _remotePort(_socket.remote_endpoint().port()), _remoteAddress(_socket.remote_endpoint().address()), _closed(false), _closing(false), _isWritingAsync(false)
    {
        _readBuffer.Resize(4096);
    }
    virtual ~Socket()
    {
        _socket.close();
    }

    Socket(Socket const& other) = delete;
    Socket(Socket&& other) = delete;
    Socket& operator=(Socket const& other) = delete;
    Socket& operator=(Socket&& other) = delete;

public:
    // virtual function
    virtual bool Update()
    {
        if(_closed) return false;

        if(_isWritingAsync || (_writeQueue.empty() && !_closing))
        {
            return true;
        }
        // handle write queue，如果是_closing，赶紧处理掉
        for(; HandleWriteQueue(); )
            ;
        return true;
    }

    virtual void OnClosed() {}

    virtual bool HandleRead() = 0;

    virtual void Start()
    {
        AsyncRead();
    }
public:
    bool IsOpen() { return !_closed && !_closing; }

    void CloseSocket()
    {
        if(_closed.exchange(true)) return;

        // 关闭本端的发送，但还可以接收对端消息，避免遗漏数据，在析构函数里再调用close全关闭
        boost::system::error_code shutdownError;
        _socket.shutdown(boost::asio::socket_base::shutdown_send, shutdownError);
        if (shutdownError)
        {
            //TC_LOG_DEBUG("network", "Socket::CloseSocket: {} errored when shutting down socket: {} ({})", GetRemoteIpAddress().to_string(),
            //    shutdownError.value(), shutdownError.message());
        }

        OnClosed();
    }

    // 上层在断开前，需要发送数据告知客户端，再关闭，可调用此接口
    void DelayCloseSocket()
    {
        if(_closing.exchange(true)) return;

        if(_writeQueue.empty())
        {
            CloseSocket();
        }
    }

    void WriteMessage(MessageBuffer&& message)
    {
        // 不加锁，因为不知道上层是多线程还是单线程，锁应该由上层去处理
        _writeQueue.push(std::move(message));
    }

    std::string GetRemoteAddress()
    {
        return _remoteAddress.to_string();
    }

    MessageBuffer& GetReadBuffer()
    {
        return _readBuffer;
    }

private:
    // 当前发送缓冲区已经满了，或者发送过程中遇到中断，则异步重新发送
    bool AsyncProcessQueue()
    {
        if(_isWritingAsync) return false;

        _isWritingAsync = true;

        _socket.async_write_some(boost::asio::null_buffers(),
            [self = this->shared_from_this()](boost::system::error_code const& error, std::size_t transferedBytes)
            {
                self->AsyncProcessQueueCallback(error, transferedBytes);
            });

        return true;
    }

    void AsyncProcessQueueCallback(boost::system::error_code const& err, std::size_t transferedBytes)
    {
        _isWritingAsync = false;
        HandleWriteQueue();
    }

    bool HandleWriteQueue()
    {
        if(_writeQueue.empty())
        {
            return false;
        }

        MessageBuffer& queuedMessage = _writeQueue.front();

        std::size_t bytesToSend = queuedMessage.GetActiveSize();

        boost::system::error_code error;
        std::size_t bytesSent = _socket.write_some(boost::asio::buffer(queuedMessage.GetReadPointer(), bytesToSend), error);

        if (error)
        {
            if (error == boost::asio::error::would_block || error == boost::asio::error::try_again)
                return AsyncProcessQueue();

            _writeQueue.pop();
            if (_closing && _writeQueue.empty())
                CloseSocket();
            return false;
        }
        else if (bytesSent == 0)    // 理论来说是对端断开了，直接CloseSocket也行
        {
            _writeQueue.pop();
            if (_closing && _writeQueue.empty())
                CloseSocket();
            return false;
        }
        else if (bytesSent < bytesToSend) // now n > 0，只发了部分，异步继续发
        {
            queuedMessage.ReadCompleted(bytesSent);
            return AsyncProcessQueue();
        }

        // 发完一个完整的消息
        _writeQueue.pop();

        if (_closing && _writeQueue.empty())
            CloseSocket();
        return !_writeQueue.empty();
    }

    void AsyncRead()
    {
        if(!IsOpen()) return;

        _readBuffer.Normalize();
        _readBuffer.EnsureFreeSpace();
        // 必须要保证回调函数被调用之前，对象是有效的，所以使用shared_from_this
        _socket.async_read_some(boost::asio::buffer(_readBuffer.GetWritePointer(), _readBuffer.GetRemainingSpace()),
            [self = this->shared_from_this()](boost::system::error_code const& error, size_t transferredBytes)
            {
                self->AsyncReadCallback(error, transferredBytes);
            });

    }

    void AsyncReadCallback(boost::system::error_code const& err, size_t transferredBytes)
    {
        if(err)
        {
            CloseSocket();
            return ;
        }
        _readBuffer.WriteCompleted(transferredBytes);
        // 已经将数据接收到了缓冲区，调用处理数据的函数
        if(!HandleRead())
        {
            CloseSocket();
            return;
        }
        // 处理完了，再次异步读取
        AsyncRead();
    }
private:
    tcp::socket _socket;
    std::atomic<bool> _closed;
    std::atomic<bool> _closing;
    bool _isWritingAsync;
    MessageBuffer _readBuffer;
    std::queue<MessageBuffer> _writeQueue;
    boost::asio::ip::address _remoteAddress;
    uint16_t _remotePort;
};

#endif